#include <QApplication>
#include <QWidget>
#include <QVBoxLayout>
#include <QLineEdit>
#include <QTextEdit>
#include <QPushButton>
#include <QComboBox>
#include <QLabel>
#include <QString>
#include <QTcpSocket>
#include <QUdpSocket>
#include <cryptlib.h>
#include <osrng.h>
#include <aes.h>
#include <modes.h>
#include <filters.h>
#include <base64.h>
using namespace CryptoPP;
using namespace std;
class Communicator {
private:
    SecByteBlock key;
    SecByteBlock iv;
    string nax(const string& plaintext) {
        try {
            AutoSeededRandomPool rng;
            rng.GenerateBlock(iv, AES::BLOCKSIZE);

            CBC_Mode<AES>::Encryption encryptor;
            encryptor.SetKeyWithIV(key, key.size(), iv);

            string ciphertext;
            StringSource(plaintext, true,
                new StreamTransformationFilter(encryptor, new StringSink(ciphertext)));

            string encoded;
            StringSource(iv + ciphertext, true, new Base64Encoder(new StringSink(encoded), false));
            return encoded;
        } catch (const Exception& e) {
            return "Encryption failed.";
        }
    }
    string fax(const string& encryptedMessage) {
        try {
            string decoded, iv_local, ciphertext;
            StringSource(encryptedMessage, true,
                new Base64Decoder(new StringSink(decoded)));

            iv_local = decoded.substr(0, AES::BLOCKSIZE);
            ciphertext = decoded.substr(AES::BLOCKSIZE);

            CBC_Mode<AES>::Decryption decryptor;
            decryptor.SetKeyWithIV(key, key.size(), (const byte*)iv_local.data());

            string plaintext;
            StringSource(ciphertext, true,
                new StreamTransformationFilter(decryptor, new StringSink(plaintext)));

            return plaintext;
        } catch (const Exception& e) {
            return "Decryption failed.";
        }
    }
public:
    Communicator() : key(AES::DEFAULT_KEYLENGTH), iv(AES::BLOCKSIZE) {
        AutoSeededRandomPool rng;
        rng.GenerateBlock(key, key.size());
    }

    string encrypt(const string& message) { return nax(message); }
    string decrypt(const string& message) { return fax(message); }
};
class ChatWindow : public QWidget {
    Q_OBJECT
private:
    QTextEdit* chatLog;
    QLineEdit* inputField;
    QComboBox* modeSelector;
    QPushButton* sendButton;
    QLineEdit* interfaceField;
    QLabel* userLabel;
    Communicator communicator;
    QTcpSocket* tcpSocket;
    QUdpSocket* udpSocket;
public:
    ChatWindow(const QString& name, QWidget* parent = nullptr) : QWidget(parent), communicator() {
        QVBoxLayout* layout = new QVBoxLayout(this);
        modeSelector = new QComboBox(this);
        modeSelector->addItem("互联网模式");
        modeSelector->addItem("物联网模式");
        layout->addWidget(modeSelector);
        interfaceField = new QLineEdit(this);
        interfaceField->setPlaceholderText("输入服务器地址或设备 ID...");
        layout->addWidget(interfaceField);
        chatLog = new QTextEdit(this);
        chatLog->setReadOnly(true);
        chatLog->setStyleSheet("background-color: black; color: green; font-family: monospace;");
        layout->addWidget(chatLog);
        inputField = new QLineEdit(this);
        inputField->setPlaceholderText("输入你的消息...");
        inputField->setStyleSheet("background-color: black; color: green; font-family: monospace;");
        layout->addWidget(inputField);
        sendButton = new QPushButton("发送", this);
        sendButton->setStyleSheet("background-color: green; color: black; font-weight: bold;");
        layout->addWidget(sendButton);
        connect(sendButton, &QPushButton::clicked, this, &ChatWindow::sendMessage);
        tcpSocket = new QTcpSocket(this);
        udpSocket = new QUdpSocket(this);
    }
private slots:
    void sendMessage() 
    {
        QString input = inputField->text();
        if (input.isEmpty()) return;
        string plaintext = input.toStdString();
        string encrypted = communicator.encrypt(plaintext);
        string decrypted = communicator.decrypt(encrypted);
        if (modeSelector->currentIndex() == 0) {
            connectToInternet(interfaceField->text(), encrypted);
        } 
        else 
        {
            connectToIoT(interfaceField->text(), encrypted);
        }

        chatLog->append("你 >> " + QString::fromStdString(plaintext));
        chatLog->append("加密 >> " + QString::fromStdString(encrypted));
        chatLog->append("解密 >> " + QString::fromStdString(decrypted));
        inputField->clear();
    }
    void connectToInternet(const QString& serverAddress, const string& message)
     {
        tcpSocket->connectToHost(serverAddress, 12345);
        if (tcpSocket->waitForConnected(1000)) {
            tcpSocket->write(message.c_str());
            tcpSocket->flush();
        }
    }

    void connectToIoT(const QString& deviceID, const string& message) {
        udpSocket->writeDatagram(message.c_str(), QHostAddress(deviceID), 54321);
    }
};
int main(int argc, char* argv[]) {
    QApplication app(argc, argv);
    ChatWindow window("Anonymous");
    window.setWindowTitle("加密聊天通讯器");
    window.resize(600, 400);
    window.show();
    return app.exec();
}
